import torch
from torch.utils.data import Dataset
from torchvision.transforms import transforms
import cv2
import numpy as np
import random
import os
from copy import deepcopy

class MyDataset(Dataset):
    def __init__(self, img_rootPath):
        super().__init__()

        data_dir = list(os.listdir(img_rootPath))

        self.img_path_arr = []
        for img_name in data_dir:
            self.img_path_arr.append(os.path.join(img_rootPath, img_name))

        self.size = 640
        self.max_mask_size = int(self.size/2)
        self.min_mask_size = 100
        self.max_timesteps = 1000

        self.to_tensor = transforms.ToTensor()

    # 正向扩散过程(加蒙板)
    def forward_diffusion(self, x0, mask, t):
        """
        x0: 原始图像 [B,C,H,W]
        t: 时间步 [B]
        beta_schedule: 噪声调度 [T]
        """
        # batch_size = x0.shape[0]
        beta_schedule = torch.tensor([1/self.max_timesteps]*self.max_timesteps)
        sqrt_alpha_bar = torch.sqrt(1. - beta_schedule).cumprod(dim=0)[t]
        sqrt_one_minus_alpha_bar = torch.sqrt(1. - sqrt_alpha_bar ** 2)
        
        # 重参数化
        epsilon = torch.randn_like(x0)
        d_mask = torch.ones_like(mask)
        d_mask[mask==1] = 0
        x_t = d_mask*x0 + (sqrt_alpha_bar.view(1, 1, 1) * x0 + sqrt_one_minus_alpha_bar.view(1, 1, 1) * epsilon)*mask
        
        return x_t, epsilon*mask
    
    def __getitem__(self, item):

        # 时间步
        timesteps = torch.tensor(random.randint(0, self.max_timesteps-1))

        # 读取图片
        img_path = self.img_path_arr[item]
        img = cv2.imread(img_path)
        img = cv2.resize(img, [self.size, self.size])

        # 创建蒙板
        rand_x = random.randint(0, self.max_mask_size-1)
        rand_y = random.randint(0, self.max_mask_size-1)
        rand_w = random.randint(self.min_mask_size, self.max_mask_size)
        rand_h = random.randint(self.min_mask_size, self.max_mask_size)
        mask = np.zeros([img.shape[0], img.shape[1], 1], dtype=np.uint8)
        mask[rand_y:rand_y+rand_h, rand_x:rand_x+rand_w, :] = 255

        # 转Tensor并采样
        x0_img = self.to_tensor(img)
        mask = self.to_tensor(mask)
        xt_img, target_epsilon = self.forward_diffusion(x0_img, mask, timesteps)        

        return xt_img, mask, timesteps, target_epsilon

    def __len__(self):
        return len(self.img_path_arr)


if __name__ == '__main__':
    dataset = MyDataset(r'D:\VOCtrainval_11-May-2012\JPEGImages')
    print(dataset[5][0])